import numpy as np
import pandas as pd
from pyswarm import pso
from model import get_flow


# 读取 limit.csv 文件
def read_limits(filename="limit.csv"):
    # 读取 limit.csv 文件
    limit_df = pd.read_csv("./data/origin/limit.csv", index_col=0)

    # 提取变量的名称、下限、上限
    names = limit_df.columns  # 第一列是索引，列名从第二列开始
    lower_limits = limit_df.loc["floor"].values  # 获取下限
    upper_limits = limit_df.loc["ceiling"].values  # 获取上限
    return names, lower_limits, upper_limits


# 适应度函数，粒子需要优化的目标函数
def fitness(coefficient):
    flow, difference = get_flow(coefficient)

    # 惩罚条件1：flow的所有元素都要大于0
    penalty_flow = np.sum(flow <= 0)  # 如果flow有元素小于等于0，惩罚增加

    # 惩罚条件2：difference要尽量接近0
    penalty_difference = (difference - 0) ** 2  # 惩罚接近0的平方

    # 适应度函数 = 惩罚项1 + 惩罚项2
    fitness_value = penalty_flow + penalty_difference
    return fitness_value


# 将满足条件的解与名称存入CSV文件
def save_results(coefficient, names, filename="optimized_results.csv"):
    data = {names[i]: coefficient[i] for i in range(len(coefficient))}
    df = pd.DataFrame([data])
    df.to_csv(
        filename, mode="a", header=not pd.io.common.file_exists(filename), index=False
    )


# 粒子群优化的主函数
def optimize():
    # 读取限制条件
    names, lower_limits, upper_limits = read_limits()

    # 粒子群优化
    lb = lower_limits  # 下限
    ub = upper_limits  # 上限

    # 运行PSO
    best_position, best_value = pso(fitness, lb, ub, swarmsize=100, maxiter=50)

    # 保存满足条件的解
    save_results(best_position, names)


# 执行优化
if __name__ == "__main__":
    optimize()
